package problem;

import datastruct.tree.Trie;

import java.util.HashMap;
import java.util.Map;

/**
 * @author RunningShrimp
 * @date 2021/5/28  11:28
 * @see <a href="https://leetcode-cn.com/problems/implement-trie-prefix-tree/">208. 实现 Trie (前缀树)</a>
 */
public class ImplementTriePrefixTree {
    public class Trie {

        private final Node root;

        private int size;

        public Trie() {
            root = new Node();
            size = 0;
        }

        public int size() {
            return size;
        }

        public void add(String word) {
            Node cur = root;

            for (int i = 0; i < word.length(); i++) {
                char c = word.charAt(i);
                if (cur.next.get(c) == null) {
                    cur.next.put(c, new Node());

                }
                cur = cur.next.get(c);
            }
            if (!cur.isWord) {
                cur.isWord = true;
                size++;
            }
        }

        public boolean contains(String word) {
            Node cur = root;

            for (int i = 0; i < word.length(); i++) {
                char c = word.charAt(i);
                if (cur.next.get(c) == null) {
                    return false;
                }
                cur = cur.next.get(c);
            }
            return cur.isWord;
        }

        public boolean isPrefix(String prefix) {
            Node cur = root;

            for (int i = 0; i < prefix.length(); i++) {
                char c = prefix.charAt(i);
                if (cur.next.get(c) == null) {
                    return false;
                }
                cur = cur.next.get(c);
            }
            return true;
        }


        private  class Node {
            boolean isWord;
            Map<Character, Node> next;

            public Node(boolean isWord) {
                this.isWord = isWord;
                next = new HashMap<>();
            }

            public Node() {
                this(false);
            }

        }
    }
}
